#include <iostream>
#include <fstream>
#include <ctime>
#include <cstdio>
#include <vector>
#include "opencv2/myopencv.hpp"

using namespace std;

int lb=6000, ub=7500;
double rg=ub-lb;

char* filename;
FILE* fin;
long long startFilePos;
long long fileSize=0;
int w, h;
int frameNum;
int curFrame;
int frameSize;
int seconds;
short* frame;
CvMat* mat;
CvMat* newmat;
CvMat* backup;
bool isSkipTo = true;

int copyStart = -1;
bool pause = false;
vector<CvKalman*> kalmans;

void loadConfig( std::string filename = "config.xml" )
{
	cv::FileStorage fs(filename, cv::FileStorage::READ);
	if (fs.isOpened() == false) return;
	fs["LowerBound"] >> lb;
	fs["UpperBound"] >> ub;
	rg = ub - lb;
}

void recovery(void)
{
	short* matPtr;
	short* backupPtr;
	matPtr = mat->data.s;
	backupPtr = backup->data.s;
	for (int i=0; i<mat->rows; ++i)
		for (int j=0; j<mat->cols; ++j, ++matPtr, ++backupPtr)
			if (*matPtr < 0)
				*matPtr = *backupPtr;
}

void kalmanRecovery(void) {
	int rowStep = 1, colStep = 1;
	int lowerB = 200, upperB = 360;
	CvMat* z_k = cvCreateMatHeader(rowStep, colStep, CV_16SC1);
	CvMat* z_k_fl = cvCreateMat(rowStep*colStep, 1, CV_32FC1);

	if (kalmans.size() == 0) {
		for (int i=lowerB; i<upperB; i+=rowStep) {
			for (int j=lowerB; j<upperB; j+=colStep) {
				CvKalman* kalman = cvCreateKalman(rowStep*colStep, rowStep*colStep);
				cvGetSubRect(mat, z_k, cvRect(j, i, colStep, rowStep));
				for (int k=0; k<rowStep*colStep; ++k)
					kalman->state_post->data.fl[k] = z_k->data.s[k];
				cvSetIdentity( kalman->measurement_matrix, cvRealScalar(1) );
				cvSetIdentity( kalman->process_noise_cov, cvRealScalar(1e-5) );
				cvSetIdentity( kalman->measurement_noise_cov, cvRealScalar(1e2) );
				cvSetIdentity( kalman->error_cov_post, cvRealScalar(1e0) );
				kalmans.push_back(kalman);
			}
		}
	}

	vector<CvKalman*>::iterator iter = kalmans.begin();
	for (int i=lowerB; i<upperB; i+=rowStep) {
		for (int j=lowerB; j<upperB; j+=colStep, ++iter) {
			CvKalman* kalman = *iter;
			cvGetSubRect(mat, z_k, cvRect(j, i, colStep, rowStep));

			if (cvGetReal2D(z_k, 0, 0) < 0) {
				const CvMat* y_k = cvKalmanPredict(kalman);
				for (int k=0; k<rowStep*colStep; ++k)
					z_k->data.s[k] = y_k->data.fl[k];
			}

		}
	}

	cvSmooth(mat, mat);

	iter = kalmans.begin();
	for (int i=lowerB; i<upperB; i+=rowStep) {
		for (int j=lowerB; j<upperB; j+=colStep, ++iter) {
			CvKalman* kalman = *iter;
			cvGetSubRect(mat, z_k, cvRect(j, i, colStep, rowStep));
			for (int k=0; k<rowStep*colStep; ++k)
				z_k_fl->data.fl[k] = z_k->data.s[k];
			cvKalmanCorrect(kalman, z_k_fl);
		}
	}

	cvReleaseMat(&z_k_fl);
	cvReleaseMatHeader(&z_k);

	cvRectangle(mat, cvPoint(lowerB,lowerB), cvPoint(upperB,upperB), cvScalar(6000,6000,6000));
}

void showFrame(void) {
	mat->data.s = frame;
	cvFlip(mat, NULL, 1);
	cvShowImage("original data", mat);

	cvConvertScale(mat, newmat, 1, -lb);
	cvConvertScale(newmat, newmat, 1/rg, 0);
	cvShowImage("original filtered data", newmat);

	recovery();
// 	kalmanRecovery();
	cvConvertScale(mat, newmat, 1, -lb);
	cvConvertScale(newmat, newmat, 1/rg, 0);
	cvShowImage("recovery filtered data", newmat);
// 	cvCopy(mat, backup);
	cvWaitKey(1000/60);
}

void skipTo(int frameNo) {
	if (isSkipTo == false) return;
	curFrame = frameNo;
	long long tmp = curFrame;
	tmp *= frameSize;
	tmp += startFilePos;
	_fseeki64(fin, tmp, SEEK_SET);
	if (pause) {
		fread_s(frame, frameSize, sizeof(short), w*h, fin);
		showFrame();
	}
}

void skipToCallback(int val) {
	skipTo(val*30);
}

void setStartCallback(int state, void* userdata) {
	copyStart = curFrame;
}

void pauseCallback(int state, void* userdata) {
	pause = !pause;
}

void prevFrameCallback(int state, void* userdata) {
	skipTo(curFrame-1);
}

void nextFrameCallback(int state, void* userdata) {
	skipTo(curFrame+1);
}

void copyCallback(int state, void* userdata) {
	FILE* fout;
	fopen_s(&fout, "copied.txt", "wb");
	fprintf_s(fout, "%d %d\n", w, h);
	_fseeki64(fin, startFilePos+copyStart*frameSize, SEEK_SET);
	int COPY_FRAMES_EACH_ITERATION = 30;
	short *newFrameBuf = new short [w*h*COPY_FRAMES_EACH_ITERATION];
	for (int i=copyStart; i<curFrame; i+=COPY_FRAMES_EACH_ITERATION) {
		int len = min(curFrame-i, COPY_FRAMES_EACH_ITERATION);
		fread_s(newFrameBuf, frameSize*COPY_FRAMES_EACH_ITERATION, frameSize, len, fin);
		fwrite(newFrameBuf, frameSize, len, fout);
	}
	fclose(fout);
}

void tryStream(void) {
	ifstream fin2(filename, ifstream::in | ifstream::binary);
	fin2.seekg(0, ios_base::end);
	cout << fin2.tellg() << endl;
	fin2.close();
}


int main(int argc, char* argv[])
{
	double cpufreq = cvGetTickFrequency();

	loadConfig();
	filename = "cloudy_3_depth.txt";
	if (argc == 2) {
		filename = argv[1];
	}

	fopen_s(&fin, filename, "rb");
	fscanf_s(fin, "%d %d\n", &w, &h);
	startFilePos = _ftelli64(fin);
	_fseeki64(fin, 0, SEEK_END);
	fileSize = _ftelli64(fin);
	_fseeki64(fin, startFilePos, SEEK_SET);

	frameSize = w*h*sizeof(short);
	frameNum = (fileSize-startFilePos) / frameSize;
	seconds = frameNum / 30;

	frame = new short[w*h];
	mat = cvCreateMatHeader(h, w, CV_16SC1);
	newmat = cvCreateMat(h, w, CV_64FC1);
	backup = cvCreateMat(h, w, CV_16SC1);
	cvZero(backup);

	int val;
	cvNamedWindow("original data");
	cvCreateTrackbar("trackbar", "original data", &val, seconds, skipToCallback);

	cvCreateButton("setStart", setStartCallback, NULL, CV_PUSH_BUTTON);
	cvCreateButton("pause", pauseCallback, NULL, CV_CHECKBOX);
	cvCreateButton("copy", copyCallback, NULL, CV_PUSH_BUTTON);
	cvCreateButton("previousFrame", prevFrameCallback, NULL, CV_PUSH_BUTTON);
	cvCreateButton("nextFrame", nextFrameCallback, NULL, CV_PUSH_BUTTON);

	curFrame = 0;
	int readFrames = 0;
	long long startTime;
	while (true) {
		if (readFrames == 0) {
			startTime = cvGetTickCount();	
		}
		if (curFrame % 30 == 0) {
			isSkipTo = false;
			cvSetTrackbarPos("trackbar", "original data", curFrame/30);
			isSkipTo = true;
		}
		fread_s(frame, frameSize, sizeof(short), w*h, fin);
		++readFrames;
		if (readFrames >= 10) {
			cout << "FPS: " << readFrames / ((cvGetTickCount() - startTime) / cpufreq / 1000000) << endl;
			readFrames = 0;
		}
		while (pause) {
			cvWaitKey(1000/30);
		}
		showFrame();

		++curFrame;
		if (feof(fin)) {
			skipTo(0);
		}
	}
	cvDestroyAllWindows();

	return 0;
}
